from __future__ import annotations

from typing import Any, Dict, List
from sqlalchemy import create_engine, text
from sqlalchemy.engine import Engine, Result
from .config import settings


def _build_target_url() -> str:
	# 使用配置的连接信息
	db_user = settings.target_db_user
	db_password = settings.target_db_password
	db_host = settings.target_db_host
	db_port = settings.target_db_port
	db_name = settings.target_db_name
	# 推荐使用 psycopg (psycopg3) 驱动
	url = f"postgresql+psycopg://{db_user}:{db_password}@{db_host}:{db_port}/{db_name}"
	print(f"Debug: PostgreSQL connection URL: {url.replace(db_password, '******')}")
	return url


_engine: Engine | None = None


def get_engine() -> Engine:
	global _engine
	if _engine is None:
		url = _build_target_url()
		# 设置 search_path 到包含短横线的 schema，需要通过 options 传入
		# 同时设置连接超时到 30 分钟（其余 read/write 超时不适用于 psycopg）
		_engine = create_engine(
			url,
			pool_pre_ping=True,
			connect_args={
				"options": f"-c search_path=\"{settings.target_db_schema}\"",
				"connect_timeout": 1800,
			},
		)
	return _engine


class DatabaseUtils:
	@staticmethod
	def execute_target_query(query: str, params: Dict[str, Any] | None = None) -> Result:
		engine = get_engine()
		# 统一在执行层做一次 PostgreSQL 反引号兼容转换，兜底所有调用方
		if (settings.target_db_type or "").lower() == "postgresql" and "`" in (query or ""):
			query = query.replace("`", '"')
		with engine.connect() as conn:
			return conn.execute(text(query), params or {})

	@staticmethod
	def get_target_tables() -> List[str]:
		t = settings.target_db_type.lower()
		if t == "mysql":
			sql = (
				"SELECT table_name FROM information_schema.tables WHERE table_schema=:db"
			)
			res = DatabaseUtils.execute_target_query(sql, {"db": settings.target_db_name})
			return [r[0] for r in res.fetchall()]
		if t == "postgresql":
			# 先查当前 search_path 下的表；若为空，再回退到 public
			sql = "SELECT tablename FROM pg_tables WHERE schemaname = current_schema()"
			res = DatabaseUtils.execute_target_query(sql)
			tables = [r[0] for r in res.fetchall()]
			if not tables:
				res = DatabaseUtils.execute_target_query(
					"SELECT tablename FROM pg_tables WHERE schemaname = :schema",
					{"schema": settings.target_db_schema or "public"}
				)
				tables = [r[0] for r in res.fetchall()]
				if not tables and (settings.target_db_schema or "public") != "public":
					res = DatabaseUtils.execute_target_query(
						"SELECT tablename FROM pg_tables WHERE schemaname = 'public'"
					)
					tables = [r[0] for r in res.fetchall()]
			return tables
		if t == "sqlite":
			sql = "SELECT name FROM sqlite_master WHERE type='table'"
			res = DatabaseUtils.execute_target_query(sql)
			return [r[0] for r in res.fetchall()]
		return []

	@staticmethod
	def get_target_table_schema(table_name: str) -> Dict[str, Any]:
		t = settings.target_db_type.lower()
		if t == "mysql":
			sql = (
				"SELECT column_name, data_type, is_nullable, column_key FROM information_schema.columns "
				"WHERE table_schema=:db AND table_name=:table ORDER BY ordinal_position"
			)
			res = DatabaseUtils.execute_target_query(sql, {"db": settings.target_db_name, "table": table_name})
			return {"columns": [dict(r._mapping) for r in res.fetchall()]}
		if t == "postgresql":
			sql = (
				"SELECT column_name, data_type, is_nullable FROM information_schema.columns "
				"WHERE table_schema = :schema AND table_name=:t ORDER BY ordinal_position"
			)
			res = DatabaseUtils.execute_target_query(sql, {"t": table_name, "schema": settings.target_db_schema or "public"})
			cols = [dict(r._mapping) for r in res.fetchall()]
			if not cols and (settings.target_db_schema or "public") != "public":
				# 回退 public
				res = DatabaseUtils.execute_target_query(sql, {"t": table_name, "schema": "public"})
				cols = [dict(r._mapping) for r in res.fetchall()]
			return {"columns": cols}
		if t == "sqlite":
			sql = f"PRAGMA table_info('{table_name}')"
			res = DatabaseUtils.execute_target_query(sql)
			return {"columns": [dict(r._mapping) for r in res.fetchall()]}
		return {"columns": []}

	@staticmethod
	async def init_database() -> None:
		# Minimal: just ensure engine can be created
		get_engine()

	@staticmethod
	def get_table_samples(table_name: str) -> List[Dict[str, Any]]:
		"""Return limited sample rows from the given table to avoid overwhelming the LLM."""
		engine = get_engine()
		# 限制每表最多返回10行
		# PostgreSQL 使用双引号引用标识符；search_path 已设置到 wd-tv-data
		sql = f'SELECT * FROM "{table_name}" LIMIT 10'
		try:
			res = DatabaseUtils.execute_target_query(sql)
			rows = [dict(r._mapping) for r in res.fetchall()]
			print(f"Debug: Retrieved {len(rows)} rows from table {table_name}")
			return rows
		except Exception as e:
			print(f"Error getting samples from {table_name}: {e}")
			return []

	@staticmethod
	def get_column_comments_map() -> Dict[str, Dict[str, str]]:
		"""Return mapping: {table_name: {column_name: comment}} for the configured DB.
		Only supports MySQL and PostgreSQL; SQLite returns empty.
		"""
		result: Dict[str, Dict[str, str]] = {}
		t = (settings.target_db_type or "").lower()
		if t == "mysql":
			sql = (
				"SELECT table_name, column_name, column_comment FROM information_schema.columns "
				"WHERE table_schema = :db"
			)
			rows = DatabaseUtils.execute_target_query(sql, {"db": settings.target_db_name}).fetchall()
			for r in rows:
				table = r[0]
				col = r[1]
				comment = r[2] or ""
				result.setdefault(table, {})[col] = comment
			return result
		if t == "postgresql":
			sql = (
				"SELECT c.relname AS table_name, a.attname AS column_name, d.description AS comment "
				"FROM pg_class c "
				"JOIN pg_namespace n ON n.oid = c.relnamespace "
				"JOIN pg_attribute a ON a.attrelid = c.oid AND a.attnum > 0 AND NOT a.attisdropped "
				"LEFT JOIN pg_description d ON d.objoid = a.attrelid AND d.objsubid = a.attnum "
				"WHERE n.nspname = :schema AND c.relkind = 'r'"
			)
			rows = DatabaseUtils.execute_target_query(sql, {"schema": settings.target_db_schema or "public"}).fetchall()
			for r in rows:
				table = r[0]
				col = r[1]
				comment = r[2] or ""
				result.setdefault(table, {})[col] = comment
			return result
		return result


